#light
#load "primitives.fs"
#load "problems.fs"
#load "basic2dplot.fs"

open LambdaMop.Problems
open LambdaMop.Primitives
open LambdaMop.Plotting
open System.Collections.Generic

type Spea2Args = {file: int; 
                  population: int} 


let spea2Selection (mop:MOP<'a,'b>) (size:int) (solutions:#seq<'a>) = 
    let eval sol = invMap sol mop.functions 
    let lenght = Seq.length solutions
    // See PISA release notes about SPEA2 
    // You can use: let k = int (sqrt ( float (args.file + args.population))) 
    let k = 1     
    let (fitness,distances) =

        // Spea strenght (number of solutions that each solution dominates)
        let strenght =
            let result = Array.create lenght 0. 
            let setStrenght i sol1 j sol2 = 
                if dominatedCompare (eval sol1) (eval sol2) = -1 then 
                    result.[i] <- result.[i] + 1.
            crossIteri setStrenght solutions solutions 
            result

        // Spea raw fitness (sum of the strenghts of every dominated solution)  
        let rawFitness =
            let result = Array.create lenght 0. 
            let setFitness i sol1 j sol2 = 
                if dominatedCompare (eval sol1) (eval sol2) = 1 then 
                    result.[i] <- result.[i] + strenght.[j]
            crossIteri setFitness solutions solutions 
            result

        // Gets the kDistance and the matrix of distances (useful for 
        // the truncation process)
        let (kDistance,distances) =
            let result = Array.create_matrix lenght lenght 0. //if only Array2 suported getting a subarray :'( 
            let setDistance i sol1 j sol2 =
                result.[j].[i] <- distance sol1 sol2
            crossIteri setDistance solutions solutions //Is not necessary to calculate the distance twice... ;-) 
            Array.iter (Array.sort compare) result  
            result |> Array.map (fun x -> 1./ (x.[k] + 2.)),result
        (Array.map2 (+) rawFitness kDistance, distances)
        
    // Returns the given number of solutions
    let envSelection size =
        let (dominated:List<int * 'a>) = new List<int * 'a>()
        let (nonDominated:List<int * 'a>) = new List<int * 'a>()
        // We build a tuple to save the index info. That way we can access in O(1) 
        // time to the already calculated Fitness and Distance. 
        let tmp i s = 
            if fitness.[i] < 1. then 
                nonDominated.Add((i,s)) 
            else
                dominated.Add((i,s));
        Seq.iteri tmp solutions
        // Get the nonDominated, if is bigger truncate, if is smaller add the best
        // Dominated solutions 
        match nonDominated.Count with
            | x when x = size -> nonDominated
            | x when x < size -> 
                dominated.Sort(compareWith (fun (x,_) -> fitness.[x]))
                nonDominated.AddRange(Seq.take (size - x) dominated)
                nonDominated
            | x -> 
                nonDominated.Sort(compareWith (fun (x,_) -> distances.[x]))
                nonDominated.GetRange(x-size, size)
    //fitness,envSelection size
    envSelection size |> Seq.map snd 


let randomWithSpeaSel (mop:MOP<'a,'b>) popSize selSize times = 
    let initial = Seq.init_finite popSize (fun _ -> mop.generator())
    let rec tmp currentIter pop =
        match currentIter with
        | 0 -> spea2Selection mop selSize pop
        | n -> 
            let selected = spea2Selection mop selSize pop
            let rest = (Seq.init_finite (popSize - selSize) (fun _ -> mop.generator())) 
            tmp (n-1) (Seq.append selected rest) 
    tmp times initial


let f1 (x: float []) = x.[0]
let f2 (x: float []) = x.[1]
let f3 (x: float []) = x.[2]

let sols = [[|-3.;3.;3.|];[|4.;3.;2.|];
            [|2.;3.;3.|];[|5.;1.;1.|];
            [|6.;0.;1.|];[|5.;5.;5.|];
            [|-1.;-1.;2.|];[|0.;0.;0.|];
            [|1.;-1.;1.|];[|0.;0.;0.|];
            [|1.;-1.;1.|];[|0.;0.;0.|]]

let pExample = seq [f1;f2;f3]

let p1 = {generator = (fun () -> [|1.;1.;1.|]); functions = pExample}
let p2 = zdt1 3 

//spea2Selection p2 3 (List.map List.of_array sols )  |> print_any
let solutions = randomWithSpeaSel p2 20 10 1000

printf "Solutions \n"

solutions |> print_any

printf "Objs \n"

let to_TupleList x =
    let to_Tuple x = 
        match x with
        | x::y::[] -> (x,y)
    List.map to_Tuple x

solutions |> List.of_seq |> List.map (fun x -> (invMap x p2.functions 
|> List.of_seq)) |> to_TupleList  
|> plot


System.Console.ReadLine()
